#!/usr/bin/env python3
# -*-coding:utf8-*-

import json
import logging
import os
import time
from pathlib import Path
from typing import Dict, List, Sequence

import numpy as np
from piper_sdk import C_PiperInterface_V2
from scipy.spatial.transform import Rotation

# from model.transforms import normalize_end, denormalize_joint, normalize_joint


log = logging.getLogger(__name__)


def arm_end_transform(status: Sequence[float], shift_len: float):
    pos, angle = status[:3], status[3:]
    # Use "x,y,z": external rotation axis. NOT "XYZ"
    Rx = Rotation.from_euler("x", angle[0], degrees=True)
    Ry = Rotation.from_euler("y", angle[1], degrees=True)
    Rz = Rotation.from_euler("z", angle[2], degrees=True)

    shift = Rz.apply(Ry.apply(Rx.apply(np.array([0, 0, shift_len]))))
    return np.concatenate([pos + shift.round(3), angle]).round(3)


# def safety_clip(xyz: Sequence[float]):
#     return np.array(xyz)
#     # x, y, z = xyz
#     # return np.array([np.clip(x, 0, 1000), y, np.clip(z, 0, 1000)])


# def to_action_val(x):
#     # mm -> um
#     # degree -> 0.001 degreee
#     return round(x * 1000)


# def from_action_val(x):
#     # um -> mm
#     # 0.001 degreee -> degree
#     return round(x / 1000, 3)


def limited_clip_joints(joints: Sequence[float]):
    j1, j2, j3, j4, j5, j6 = joints
    j1 = np.clip(j1, -150, 150)
    j2 = np.clip(j2, 0, 180)
    j3 = np.clip(j3, -170, 0)
    j4 = np.clip(j4, -100, 100)
    j5 = np.clip(j5, -70, 70)
    j6 = np.clip(j6, -120, 120)
    return np.array([j1, j2, j3, j4, j5, j6])


class PiperArm:
    """
    Unify the measurement unit:
    Length: mm
    Angle: radian
    Time: s
    Speed: mm/s
    """

    __speed: int = 40
    __response_t: float = 0.01
    __timeout: int = 5
    __ctrl_mode: int = 0x00
    __raise_height: int = 300
    __gripper_delay: float = 0.3
    __exe_delay: float = 0.5

    def __init__(
        self,
        can: str = "can0",
        end_joint_model=None,
        planning_model=None,
        # arm_joint_model=None,
        wait_execution: bool = True,
        gripper_length: float = 130,
    ):  # in mm

        # Reset the robot
        self.can = can
        self.wait_execution = wait_execution
        self.gripper_length = gripper_length
        self.arm = C_PiperInterface_V2(self.can)
        self.arm.ConnectPort()
        if self.reset():
            print(f"Sucessfully connect to the robot CAN={can}")
            print(f"Arm {can} enabled")

        self.go_zero()

        self.end_joint_model = end_joint_model
        self.planning_model = planning_model
        # self.arm_joint_model = None

        if end_joint_model is not None:
            self.end_joint_model = end_joint_model
        # if arm_joint_model is not None:
        # self.arm_joint_model = arm_joint_model

    def reset(self):
        s_time = time.time()
        while not all(self.arm.GetArmEnableStatus()):
            self.arm.EnableArm(7)
            self.go_zero()
            time.sleep(self.__response_t)
            if time.time() - s_time > self.__timeout:
                log.error("Arm reset timeout!")
                return False
        return True

    def set_speed(self, speed: int):
        self.__speed = speed

    def go_zero(self):
        # use joint mode
        self.arm.ModeCtrl(0x01, 0x01, self.__speed, self.__ctrl_mode)
        self.arm.JointCtrl(0, 0, 0, 0, 0, 0)
        if self.wait_execution:
            self.wait_motion_done()
        self.arm.GripperCtrl(0, 1000, 0x01, 0)

    def raise_and_move_to(self, xyz: Sequence[float]):
        current_xyz = self.get_end_pos()[:3]

        # can be modified
        mid_point1 = np.array(current_xyz) * 0.9 + np.array(xyz) * 0.1
        mid_point1[2] += self.__raise_height
        self.move_to(mid_point1)

        mid_point2 = np.array(current_xyz) * 0.1 + np.array(xyz) * 0.9
        mid_point2[2] += self.__raise_height
        self.move_to(mid_point2)
        time.sleep(self.__exe_delay)  # execution + time
        self.move_to(xyz)

    def planning_move_to(self, goal_xyz: Sequence[float]):
        pass
        # if self.planning_model is None:
        #     print("planning model is not defined!")
        #     return
        # # assert self.end_joint_model.is_planning, "Only support planning model"
        # current_joint = normalize_joint(self.get_arm_joint_status()['joint'])
        # current_xyz = normalize_end(self.get_end_pos())[:3]
        # goal_xyz = normalize_end((self.safety_clip(goal_xyz)))
        # # for j in self.end_joint_model(current_xyz, tar_xyz)[0]:
        # #     self.move_to_arm_joint(j)
        # cond = jnp.concat([current_joint, current_xyz, goal_xyz], axis=-1)

        # return self.move_to_arm_joint(denormalize_joint(self.planning_model(cond)[0, 0]))

        # for j in self.planning_model(cond)[0]:  # remove batch
        #     self.move_to_arm_joint(denormalize_joint(j))
        #     time.sleep(0.1)

    def move_to(self, xyz: Sequence[float]):
        """
        USE THIS TO REACH CONTACT POINT !!!!

        move to the contact position (gripper)
        xyz: the coordinates in mm
        first clip the xyz to the same zone,
        If model loaded: predict the joint state -> joint move
        If no model:
        """
        # give the axis value on the table!

        xyz = self.safety_clip(xyz)
        if self.end_joint_model is not None:
            # inference in the normed space
            pred_joint = self.end_joint_model(xyz)[0]  # remove batch
            return self.move_to_arm_joint(self.limited_clip_joints(pred_joint))
        return self.move_to_arm_pos(self.est_arm_pos_from_gripper_pos(xyz))

    def move_and_grasp(self, xyz: Sequence[float]):
        if self.from_action_val(self.arm.GetArmGripperMsgs().gripper_state.grippers_angle) < 50:
            self.open_gripper()

        self.move_to(xyz)
        self.close_gripper()

    def grasp(self):
        if self.from_action_val(self.arm.GetArmGripperMsgs().gripper_state.grippers_angle) < 50:
            self.open_gripper()
        self.close_gripper()

    def open_gripper(self, width: int = 100):
        # need some time to finish the operation
        self.wait_motion_done()
        self.arm.GripperCtrl(width * 1000, 1000, 0x01, 0)
        time.sleep(self.__gripper_delay)

    def close_gripper(self):
        # need some time to finish the operation
        self.wait_motion_done()
        self.arm.GripperCtrl(0, 1000, 0x01, 0)
        time.sleep(self.__gripper_delay)

    def move_to_arm_pos(self, pos: np.ndarray):
        """
        xyz: mm
        rpy: radian
        """
        assert len(pos) == 6, "pos=[x, y, z, rx, ry, rz]"

        xyz = np.round(pos[:3] * 1000).tolist()
        rpy = np.round(np.deg2rad(pos[3:]) * 1000).tolist()

        # if self.arm_joint_model is not None:
        #     pred_joint = self.arm_joint_model(pos)[0]
        #     return self.move_to_arm_joint(self.limited_clip_joints(pred_joint))
        # tar = [self.to_action_val(_) for _ in pos]

        # "Position" mode
        self.arm.ModeCtrl(0x01, 0x00, self.__speed, self.__ctrl_mode)
        self.arm.EndPoseCtrl(*(xyz + rpy))
        if self.wait_execution:
            return self.wait_motion_done()

    def move_to_arm_joint(self, joints: np.ndarray):
        """
        joints: radian -> 0.001 degree
        """
        assert len(joints) == 6, "Numer of target joint status is 6"
        machine_joints = np.round(np.deg2rad(joints) * 1000)

        # "Joint" mode
        self.arm.ModeCtrl(0x01, 0x01, self.__speed, self.__ctrl_mode)
        self.arm.JointCtrl(*machine_joints)
        if self.wait_execution:
            return self.wait_motion_done()

    def grad_move_to(self, pos: Sequence[float]):
        norm_j = normalize_joint(self.get_arm_joint_status()["joint"])
        norm_x = normalize_end(pos)

        nx_j = denormalize_joint(norm_j + self.fk_model.get_delta_joint(norm_j, norm_x)[0])

        # p = np.array([608.30336375, 68.6895842, 50.59737271])
        # for i in range(1000):
        #     self.grad_move_to(p)
        # print(denormalize_end(self.fk_model(normalize_joint(self.get_arm_joint_status()['joint']))))
        # print(self.get_end_pos())

        return self.move_to_arm_joint(nx_j)

    def get_arm_pos_status(
        self,
    ) -> Dict:
        arm_pos = self.arm.GetArmEndPoseMsgs().end_pose
        pos = [arm_pos.X_axis, arm_pos.Y_axis, arm_pos.Z_axis, arm_pos.RX_axis, arm_pos.RY_axis, arm_pos.RZ_axis]
        return {
            "position": np.asarray([self.from_action_val(_) for _ in pos]),
            "gripper": np.asarray(self.from_action_val(self.arm.GetArmGripperMsgs().gripper_state.grippers_angle)),
        }

    def get_arm_joint_status(
        self,
    ) -> Dict:
        arm_status = self.arm.GetArmJointMsgs().joint_state
        jo = [
            arm_status.joint_1,
            arm_status.joint_2,
            arm_status.joint_3,
            arm_status.joint_4,
            arm_status.joint_5,
            arm_status.joint_6,
        ]
        return {
            "joint": np.asarray([self.from_action_val(_) for _ in jo]),
            "gripper": np.asarray(self.from_action_val(self.arm.GetArmGripperMsgs().gripper_state.grippers_angle)),
        }

    def calibrate_gripper_length(self):
        num_points = 20
        n_ = 0
        # length_samples = []

        x, z = [], []

        while n_ < num_points:
            _ = input("(Tech mode) Contact the gripper to the table and press any key...")
            status = self.get_arm_pos_status()["position"]
            pos, angle = status[:3], status[3:]
            # Use "x,y,z": external rotation axis. NOT "XYZ"
            Rx = Rotation.from_euler("x", angle[0], degrees=True)
            Ry = Rotation.from_euler("y", angle[1], degrees=True)
            Rz = Rotation.from_euler("z", angle[2], degrees=True)
            rotation_matrx = Rz.as_matrix() @ Ry.as_matrix() @ Rx.as_matrix()
            x.append(rotation_matrx[2, 2])
            z.append(pos[2])
            # length_samples.append(-pos[2] / rotation_matrx[2, 2])
            print(f"add points {x[-1]}, {z[-1]}")
            n_ += 1

        x_hom = np.hstack([np.array(x)[:, None], np.ones((num_points, 1))])
        M, _, _, _ = np.linalg.lstsq(x_hom, np.array(z), rcond=None)
        print(f"Estimated: gripper length:{-M[0]} with matrix {M}")

        # self.gripper_length = round(sum(length_samples) / num_points, 3)
        # print(f"Reset the gipper length to:{self.gripper_length}")

    def get_end_pos(self):
        status = self.get_arm_pos_status()["position"]
        return arm_end_transform(status, self.gripper_length)
        # pos, angle = status[:3], status[3:]
        # # Use "x,y,z": external rotation axis. NOT "XYZ"
        # Rx = Rotation.from_euler("x", angle[0], degrees=True)
        # Ry = Rotation.from_euler("y", angle[1], degrees=True)
        # Rz = Rotation.from_euler("z", angle[2], degrees=True)

        # shift = Rz.apply(Ry.apply(Rx.apply(np.array([0, 0, self.__gripper_length]))))
        # return np.concat([pos + shift.round(3), angle]).round(3)

    def get_contact_gripper_angle(self, xyz: Sequence[float]):
        # Given the arm_xyz, predict the gripper status suitable for picking
        x, y, z = xyz
        joint1 = np.arctan(y / x) / np.pi * 180
        rx = (np.sign(joint1) * 180).round(3)
        ry = (np.arccos(min(z / self.gripper_length, 1)) / np.pi * 180).round(3)
        rz = (joint1 - 180 * np.sign(joint1)).round(3)
        return float(rx), float(ry), float(rz)

    def gripper_pos_to_arm_pos(self, pos: Sequence[float]):
        # if knowing the full cooridnates of the gripper, we can estimate the arm position
        x, y, z, rx, ry, rz = pos
        contact_xyz = np.array([x, y, z])
        Rx = Rotation.from_euler("x", rx, degrees=True)
        Ry = Rotation.from_euler("y", ry, degrees=True)
        Rz = Rotation.from_euler("z", rz, degrees=True)
        arm_xyz = contact_xyz - Rz.apply(Ry.apply(Rx.apply(np.array([0, 0, self.gripper_length]))))
        return np.concatenate([arm_xyz, np.array([rx, ry, rz])])

    def est_arm_pos_from_gripper_pos(self, xyz: Sequence[float], ry: float = 10):
        # ry \approx the angle from gripper's vertical postion
        x, y, z = xyz
        joint1 = np.arctan(y / x) / np.pi * 180
        rx = (np.sign(joint1) * 179).round(3)
        rz = (joint1 - 180 * np.sign(joint1)).round(3)

        ry_rad = ry / 180 * np.pi
        xy_angle = np.arctan(y / x)
        x = x - self.gripper_length * np.sin(ry_rad) * np.cos(xy_angle)
        y = y - self.gripper_length * np.sin(ry_rad) * np.sin(xy_angle)
        z = z + self.gripper_length * np.cos(ry_rad)

        return np.array([x, y, z, rx, ry, rz])

    def elliptic_move(self, start_xyz, end_xyz, num_points: int = 5):

        for p_ in self.elliptic_traj(start_xyz, end_xyz, num_points=num_points):
            self.move_to(p_)

    def raise_up(self, height: float = None):
        if height is None:
            height = self.__raise_height
        xyz = np.array([*self.get_end_pos()[:2], height])
        return self.move_to(xyz)

    def C_move(self, start_xyz, end_xyz):

        est_start_6d = self.est_arm_pos_from_gripper_pos(start_xyz)
        est_end_6d = self.est_arm_pos_from_gripper_pos(end_xyz)

        mid = (est_start_6d + est_end_6d) / 2
        mid[2] = self.__raise_height

        self.arm.MotionCtrl_2(0x01, 0x03, 30, 0x00)
        self.arm.EndPoseCtrl(*(self.to_action_val(_) for _ in est_start_6d))
        self.arm.MoveCAxisUpdateCtrl(0x01)
        time.sleep(self.__exe_delay)
        self.arm.EndPoseCtrl(*(self.to_action_val(_) for _ in mid))
        self.arm.MoveCAxisUpdateCtrl(0x02)
        time.sleep(self.__exe_delay)
        self.arm.EndPoseCtrl(*(self.to_action_val(_) for _ in est_end_6d))
        self.arm.MoveCAxisUpdateCtrl(0x03)
        time.sleep(self.__exe_delay)

    def is_in_motion(self):
        return bool(self.arm.GetArmStatus().arm_status.motion_status)

    def wait_motion_done(self):
        start_t = time.time()
        time.sleep(self.__response_t)
        while self.is_in_motion():
            time.sleep(self.__response_t)
            if time.time() - start_t > self.__timeout:
                log.error(f"Excution exceeds timeout threshold={self.__timeout}s")
                return False
        return True

    def is_gripping_sth(self):
        return self.arm.GetArmGripperMsgs().gripper_state.grippers_effort < -200

    def follow_traj(self, filename: Path, mode: str):
        with open(filename, "r") as f:
            status = json.load(f)

        if mode == "J":
            data = status["joint"]
            exe_fn = self.move_to_arm_joint
        elif mode == "P":
            data = status["arm"]
            exe_fn = self.move_to_arm_pos
        else:
            raise ValueError(f"Not supported mode={mode}")

        for p_ in data:
            exe_fn(p_)
            time.sleep(0.1)

    def record_traj(self, filename: Path = "data/json_data/trajectories.json"):
        arms, joints, ends, lengths = [], [], [], []
        try:
            while True:
                a, j, e = [], [], []
                while self.arm.GetArmStatus().arm_status.teach_status == 1:
                    a.append(self.get_arm_pos_status()["position"].tolist())
                    j.append(self.get_arm_joint_status()["joint"].tolist())
                    e.append(self.get_end_pos().tolist())
                    time.sleep(0.1)
                    print(f"Recorded {len(a)} points", end="\r")

                if a and j and e:
                    arms += a
                    joints += j
                    ends += e
                    lengths.append(len(a))
                    print(f"Recorded {len(a)} points: Done!", end="\n")
                else:
                    print("No points recorded, please move the arm to record...", end="\r")
                    time.sleep(0.1)
        except KeyboardInterrupt:
            with open(filename, "w") as f:
                json.dump(
                    {
                        "arms": arms,
                        "joints": joints,
                        "ends": ends,
                        "lengths": lengths,
                    },
                    f,
                )

    def record_points(self, save_dir: Path = "data/json_data", method: str = "auto", tag: str = ""):
        n_ = 0
        joints, arm = [], []

        if not os.path.exists(save_dir):
            os.mkdir(save_dir)

        try:
            while True:
                if method != "auto":
                    key = input("Place the gripper and press ENTER...")
                    if key == "q":
                        assert 1 == 0
                joints.append(self.get_arm_joint_status()["joint"].tolist())
                arm.append(self.get_arm_pos_status()["position"].tolist())
                # gripper.append(self.get_end_pos().tolist())
                n_ += 1

                time.sleep(0.02)
                print(f"Saved {n_} points", end="\r")

        except KeyboardInterrupt:

            f_name = os.path.join(save_dir, f"{self.can}_{tag}_{round(time.time())}.json")
            with open(f_name, "w") as f:
                json.dump(
                    {
                        "joint": joints,
                        "arm": arm,
                    },
                    f,
                )
                print(f"{n_} data points are saved to {f_name}")

    def elliptic_traj(self, start_xyz, end_xyz, num_points: int = 5):

        thetas = np.linspace(0, np.pi, num_points)
        w1, w2 = 0.5 * (1 + np.cos(thetas)), 0.5 * (1 - np.cos(thetas))

        # (num_points, 2)
        mid_points = w1[:, None] * np.array(start_xyz[:2])[None, :] + w2[:, None] * np.array(end_xyz[:2])[None, :]
        heights = self.__raise_height * np.sin(thetas) + max(start_xyz[-1], end_xyz[-1])
        return np.concatenate([mid_points, heights[:, None]], axis=-1)

    @staticmethod
    def safety_clip(xyz: Sequence[float]):
        return np.array(xyz)
        # x, y, z = xyz
        # return np.array([np.clip(x, 0, 1000), y, np.clip(z, 0, 1000)])

    @staticmethod
    def to_action_val(x):
        # mm -> um
        # degree -> 0.001 degreee
        return round(x * 1000)

    @staticmethod
    def from_action_val(x):
        # um -> mm
        # 0.001 degreee -> degree
        return round(x / 1000, 3)

    @staticmethod
    def limited_clip_joints(joints: Sequence[float]):
        j1, j2, j3, j4, j5, j6 = joints
        j1 = np.clip(j1, -150, 150)
        j2 = np.clip(j2, 0, 180)
        j3 = np.clip(j3, -170, 0)
        j4 = np.clip(j4, -100, 100)
        j5 = np.clip(j5, -70, 70)
        j6 = np.clip(j6, -120, 120)
        return np.array([j1, j2, j3, j4, j5, j6])


def test_reaching_contact_points():
    arm = PiperArm("can0")
    # arm.record_points("points.json", num_points=20)
    with open("reaching_test.json", "r") as f:
        poss = json.load(f)
    for p_ in poss["g_pos"]:
        arm.go_zero()
        if arm.move_to(p_[:3]):
            print(f"Successfully reaching {p_[:3]}: using {arm.get_end_pos()}")
        else:
            print(f"Failed reaching {p_[:3]}")
    arm.go_zero()


def point_test():
    arm = PiperArm("can0")
    # arm.record_points("points.json", num_points=20)
    with open("points.json", "r") as f:
        poss = json.load(f)

    i = 100
    arm.go_zero()

    arm.move_to_arm_joint(poss["joint"][i])
    arm.go_zero()
    arm.move_to_arm_pos(poss["a_pos"][i])
    arm.go_zero()
    arm.move_to(poss["g_pos"][i][:3])
    arm.go_zero()

    print(poss["a_pos"][:3][i])
    print(poss["g_pos"][:3][i])


if __name__ == "__main__":

    from model import *
    from model.diffusion.diffusion_ctrl import *

    # ej_model = DiffusionCtrlModel(
    #     in_dim=3,
    #     out_dim=6,
    #     model_file="diffusion_end_joint_v0.ckp",  # "diffusion_end_joint_all_normed.ckp",
    #     dropout_rate=0.1,
    #     temperature=0.1,
    #     fk_model=MLPCtrlModel(in_dim=6, out_dim=3, model_file="all_normed_FK.ckp").model,
    # )

    ej_model = DiffusionIKModel(
        in_dim=3,
        out_dim=6,
        model_file="diffusion_end_joint_T10.ckp",  #
        # dropout_rate=0.1,
        denoising_steps=10,
        # model_file="diffusion_end_joint_nodropout_v1.ckp",
        temperature=0.1,
        num_act_samples=100,
        # fk_model=EnsembleFKModel(in_dim=6, out_dim=3, num_heads=10, model_file="ensemble_FK.ckp"),
        fk_model=MlpPartialIKModel(in_dim=6, out_dim=3, model_file="all_normed_FK.ckp").model,
    )

    # plan_model = DiffusionPlanningModel(
    #     in_dim=12,
    #     out_dim=6,
    #     model_file="diffusion_planning_T10_H10_S1.ckp",  # "diffusion_end_joint_all_normed.ckp",
    #     horizon_steps=10,
    #     temperature=0.1,
    #     denoising_steps=10,
    #     num_act_samples=1,
    #     # fk_model=MLPCtrlModel(in_dim=6, out_dim=3, model_file="all_normed_FK.ckp").model,
    # )

    # self = PiperArm("can0", wait_execution=False, end_joint_model=ej_model)
    # self.planning_model = plan_model
    self = PiperArm("can0", wait_execution=False)

    self.go_zero()
    # p = np.array([435.32488388 , -24.77839234  , 1.09628048])
    # p = np.array([401.21426989, 40.37996443,  5.94443308])
    p = np.array([408.30336375, 68.6895842, 15.59737271])
    # p = np.array([ 55.59858803, -46.4059748 ,  -1.60010257]) ## danger!

    # pp = np.array([10.532, 85.783, -25.349, 2.987, 21.791, -1.417])

    # self.move_to_arm_pos(np.array([ 411.856,  -17.604,  130.043, -177.575,   11.987, -178.31 ]))
    # [ 419.21315514 -109.77677174    3.67021936]
    # self.move_to(p)
    # self.planning_move_to(p)
    # self.go_zero()
